from collections import namedtuple
from typing import Iterator

from torch import optim
from torch.nn import Parameter, Module


class OptimizerSpec:
    def __init__(self, constructor, kwargs):
        optimizer = namedtuple("OptimizerSpec", ["constructor", "kwargs"])
        self.optimizer = optimizer(constructor=constructor, kwargs=kwargs)

    def build(self, parameters: Iterator[Parameter]):
        return self.optimizer.constructor(parameters, **self.optimizer.kwargs)


def rms_prop(model: Module):
    # 学习率
    learning_rate = 0.001
    # 衰减率（alpha，用于计算梯度的滑动平均）
    alpha = 0.95
    # epsilon（eps，用于数值稳定性）
    eps = 0.01
    kwargs = dict(lr=learning_rate, alpha=alpha, eps=eps)
    return OptimizerSpec(constructor=optim.RMSprop, kwargs=kwargs).build(model.parameters())


def sdg(model: Module):
    # 学习率
    learning_rate = 0.001
    # 动量
    momentum = 0.9
    kwargs = dict(lr=learning_rate, momentum=momentum)
    return OptimizerSpec(constructor=optim.SGD, kwargs=kwargs).build(model.parameters())


if __name__ == '__main__':
    from main import DQN, BasicBlock
    dqn = DQN(in_channels=3, block=BasicBlock, num_actions=10)
    opti = rms_prop(dqn)
    print(opti)
